/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://aws.amazon.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file is distributed
 * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 * express or implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */

package software.amazon.awssdk.codegen.customization.processors;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import software.amazon.awssdk.codegen.customization.CodegenCustomizationProcessor;
import software.amazon.awssdk.codegen.internal.Utils;
import software.amazon.awssdk.codegen.model.config.customization.ModifyModelShapeModifier;
import software.amazon.awssdk.codegen.model.config.customization.ShapeModifier;
import software.amazon.awssdk.codegen.model.intermediate.EnumModel;
import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
import software.amazon.awssdk.codegen.model.intermediate.MemberModel;
import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
import software.amazon.awssdk.codegen.model.service.Member;
import software.amazon.awssdk.codegen.model.service.ServiceModel;
import software.amazon.awssdk.codegen.model.service.Shape;

/**
 * This processor handles all the modification on the shape members in the
 * pre-process step and then removes all the excluded shapes in the post-process
 * step.
 *
 * This processor can also modify the names of the enum values generated by the
 * code generator.
 */
final class ShapeModifiersProcessor implements CodegenCustomizationProcessor {

    private static final String ALL = "*";
    private final Map<String, ShapeModifier> shapeModifiers;

    ShapeModifiersProcessor(
            Map<String, ShapeModifier> shapeModifiers) {
        this.shapeModifiers = shapeModifiers;
    }

    @Override
    public void preprocess(ServiceModel serviceModel) {

        if (shapeModifiers == null) {
            return;
        }

        for (Entry<String, ShapeModifier> entry : shapeModifiers.entrySet()) {
            String key = entry.getKey();
            ShapeModifier modifier = entry.getValue();

            if (ALL.equals(key)) {
                for (Shape shape : serviceModel.getShapes().values()) {
                    preprocessModifyShapeMembers(serviceModel, shape, modifier);
                }

            } else {
                Shape shape = serviceModel.getShapes().get(key);

                if (shape == null) {
                    throw new IllegalStateException(
                            "ShapeModifier customization found for " + key
                            + ", but this shape doesn't exist in the model!");
                }

                preprocessModifyShapeMembers(serviceModel, shape, modifier);
            }
        }
    }

    @Override
    public void postprocess(IntermediateModel intermediateModel) {
        if (shapeModifiers == null) {
            return;
        }

        for (Entry<String, ShapeModifier> entry : shapeModifiers.entrySet()) {
            String key = entry.getKey();
            ShapeModifier modifier = entry.getValue();

            if (ALL.equals(key)) {
                continue;
            }

            List<ShapeModel> shapeModels = Utils.findShapesByC2jName(intermediateModel, key);
            if (shapeModels.isEmpty()) {
                throw new IllegalStateException(String.format(
                    "Cannot find c2j shape [%s] in the intermediate model when processing " +
                    "customization config shapeModifiers.%s",
                    key, key));
            }

            shapeModels.forEach(shapeModel -> {
                if (modifier.getStaxTargetDepthOffset() != null) {
                    shapeModel.getCustomization().setStaxTargetDepthOffset(modifier.getStaxTargetDepthOffset());
                }

                if (modifier.isExcludeShape()) {
                    shapeModel.getCustomization().setSkipGeneratingModelClass(true);
                    shapeModel.getCustomization().setSkipGeneratingMarshaller(true);
                    shapeModel.getCustomization().setSkipGeneratingUnmarshaller(true);
                } else if (modifier.getModify() != null) {
                    // Modifies properties of a member in shape or shape enum.
                    // This customization currently support modifying enum name
                    // and marshall/unmarshall location of a member in the Shape.

                    modifier.getModify().stream().flatMap(m -> m.entrySet().stream()).forEach(memberEntry ->
                        postprocessModifyMemberProperty(shapeModel, memberEntry.getKey(), memberEntry.getValue())
                    );
                }
            });
        }
    }

    /**
     * Override name of the enums, marshall/unmarshall location of the
     * members in the given shape model.
     */
    private void postprocessModifyMemberProperty(ShapeModel shapeModel, String memberName,
                                                 ModifyModelShapeModifier modifyModel) {
        if (modifyModel.getEmitEnumName() != null) {
            EnumModel enumModel = shapeModel.findEnumModelByValue(memberName);
            if (enumModel == null) {
                throw new IllegalStateException(
                        String.format("Cannot find enum [%s] in the intermediate model when processing "
                                      + "customization config shapeModifiers.%s", memberName, memberName));
            }
            enumModel.setName(modifyModel.getEmitEnumName());
        }

        if (modifyModel.getEmitEnumValue() != null) {
            EnumModel enumModel = shapeModel.findEnumModelByValue(memberName);
            if (enumModel == null) {
                throw new IllegalStateException(
                        String.format("Cannot find enum [%s] in the intermediate model when processing "
                                      + "customization config shapeModifiers.%s", memberName, memberName));
            }
            enumModel.setValue(modifyModel.getEmitEnumValue());
        }

        if (modifyModel.getMarshallLocationName() != null) {
            MemberModel memberModel = shapeModel.findMemberModelByC2jName(memberName);
            memberModel.getHttp().setMarshallLocationName(modifyModel.getMarshallLocationName());
        }

        if (modifyModel.getUnmarshallLocationName() != null) {
            MemberModel memberModel = shapeModel.findMemberModelByC2jName(memberName);
            memberModel.getHttp().setUnmarshallLocationName(modifyModel
                                                                    .getUnmarshallLocationName());
        }

        if (modifyModel.isIgnoreDataTypeConversionFailures()) {
            MemberModel memberModel = shapeModel.findMemberModelByC2jName(memberName);
            memberModel.ignoreDataTypeConversionFailures(true);
        }

    }

    /**
     * Exclude/modify/inject shape members
     */
    private void preprocessModifyShapeMembers(ServiceModel serviceModel, Shape shape, ShapeModifier modifier) {

        if (modifier.getModify() != null) {
            for (Map<String, ModifyModelShapeModifier> modifies : modifier.getModify()) {
                for (Entry<String, ModifyModelShapeModifier> entry : modifies.entrySet()) {

                    String memberToModify = entry.getKey();
                    ModifyModelShapeModifier modifyModel = entry.getValue();

                    doModifyShapeMembers(serviceModel, shape, memberToModify, modifyModel);
                }
            }
        }

        if (modifier.getExclude() != null) {
            for (String memberToExclude : modifier.getExclude()) {

                if (shape.getRequired() != null &&
                    shape.getRequired().contains(memberToExclude)) {
                    throw new IllegalStateException(
                            "ShapeModifier.exclude customization found for "
                            + memberToExclude
                            + ", but this member is marked as required in the model!");
                }

                if (shape.getMembers() != null) {
                    shape.getMembers().remove(memberToExclude);
                }
            }
        }

        if (modifier.getInject() != null) {
            for (Map<String, Member> injects : modifier.getInject()) {
                if (shape.getMembers() == null) {
                    shape.setMembers(new HashMap<>());
                }
                shape.getMembers().putAll(injects);
            }
        }

        if (modifier.isUnion() != null) {
            shape.setUnion(modifier.isUnion());
        }
    }

    private void doModifyShapeMembers(ServiceModel serviceModel, Shape shape, String memberToModify,
                                      ModifyModelShapeModifier modifyModel) {
        if (modifyModel.isDeprecated()) {
            Member member = shape.getMembers().get(memberToModify);
            member.setDeprecated(true);
            if (modifyModel.getDeprecatedMessage() != null) {
                member.setDeprecatedMessage(modifyModel.getDeprecatedMessage());
            }
        }
        // Currently only supports emitPropertyName which is to rename the member
        if (modifyModel.getEmitPropertyName() != null) {
            Member member = shape.getMembers().remove(memberToModify);

            // if location name is not present, set it to the original name
            // to avoid breaking marshaller code
            if (member.getLocationName() == null) {
                member.setLocationName(memberToModify);
            }

            if (modifyModel.isExistingNameDeprecated()) {
                member.setDeprecatedName(memberToModify);
            }

            shape.getMembers().put(modifyModel.getEmitPropertyName(), member);
        }
        if (modifyModel.getEmitAsType() != null) {
            // Must create a shape for the primitive type.
            Shape newShapeForType = new Shape();
            newShapeForType.setType(modifyModel.getEmitAsType());
            String shapeName = "SDK_" + modifyModel.getEmitAsType();
            serviceModel.getShapes().put(shapeName, newShapeForType);

            shape.getMembers().get(memberToModify).setShape(shapeName);
        }
    }

}
